"""
资源监控服务
监控系统资源使用情况，包括CPU、内存、GPU等
"""
import asyncio
import logging
import psutil
import time
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Any
from dataclasses import dataclass
import json

from services.ai_infrastructure.gpu_resource_manager import gpu_resource_manager
from services.ai_infrastructure.ai_task_scheduler import ai_task_scheduler
from core.config import get_settings

logger = logging.getLogger(__name__)
settings = get_settings()


@dataclass
class SystemMetrics:
    """系统指标数据类"""
    timestamp: datetime
    cpu_percent: float
    memory_percent: float
    memory_used: int  # MB
    memory_total: int  # MB
    disk_percent: float
    disk_used: int  # GB
    disk_total: int  # GB
    network_sent: int  # bytes
    network_recv: int  # bytes
    load_average: List[float]
    process_count: int
    temperature: Optional[float] = None


@dataclass
class GPUMetrics:
    """GPU指标数据类"""
    timestamp: datetime
    gpu_id: int
    utilization: float
    memory_used: int  # MB
    memory_total: int  # MB
    temperature: Optional[float]
    power_usage: Optional[float]
    fan_speed: Optional[float]


@dataclass
class AlertRule:
    """告警规则数据类"""
    rule_id: str
    name: str
    metric_type: str  # cpu, memory, gpu_utilization, gpu_temperature, etc.
    threshold: float
    operator: str  # >, <, >=, <=, ==
    duration: int  # 持续时间(秒)
    enabled: bool
    last_triggered: Optional[datetime] = None


class ResourceMonitor:
    """资源监控器"""
    
    def __init__(self):
        self.system_metrics_history: List[SystemMetrics] = []
        self.gpu_metrics_history: Dict[int, List[GPUMetrics]] = {}
        self.alert_rules: Dict[str, AlertRule] = {}
        self.active_alerts: Dict[str, Dict[str, Any]] = {}
        
        # 监控配置
        self.monitoring_interval = 30  # 30秒
        self.history_retention_hours = 24  # 保留24小时历史数据
        self.max_history_points = 2880  # 24小时 * 60分钟 * 2 (每30秒一个点)
        
        # 监控任务
        self._monitoring_task = None
        self._cleanup_task = None
        self._running = False
        
        # 网络统计基线
        self._last_network_stats = None
        
        # 初始化默认告警规则
        self._init_default_alert_rules()
    
    def _init_default_alert_rules(self):
        """初始化默认告警规则"""
        default_rules = [
            AlertRule(
                rule_id="cpu_high",
                name="CPU使用率过高",
                metric_type="cpu_percent",
                threshold=80.0,
                operator=">=",
                duration=300,  # 5分钟
                enabled=True
            ),
            AlertRule(
                rule_id="memory_high",
                name="内存使用率过高",
                metric_type="memory_percent",
                threshold=85.0,
                operator=">=",
                duration=300,
                enabled=True
            ),
            AlertRule(
                rule_id="disk_high",
                name="磁盘使用率过高",
                metric_type="disk_percent",
                threshold=90.0,
                operator=">=",
                duration=600,  # 10分钟
                enabled=True
            ),
            AlertRule(
                rule_id="gpu_utilization_high",
                name="GPU利用率过高",
                metric_type="gpu_utilization",
                threshold=95.0,
                operator=">=",
                duration=600,
                enabled=True
            ),
            AlertRule(
                rule_id="gpu_temperature_high",
                name="GPU温度过高",
                metric_type="gpu_temperature",
                threshold=85.0,
                operator=">=",
                duration=180,  # 3分钟
                enabled=True
            )
        ]
        
        for rule in default_rules:
            self.alert_rules[rule.rule_id] = rule
    
    async def start(self):
        """启动资源监控"""
        if self._running:
            return
        
        self._running = True
        self._monitoring_task = asyncio.create_task(self._monitoring_loop())
        self._cleanup_task = asyncio.create_task(self._cleanup_loop())
        
        logger.info("Resource Monitor started")
    
    async def stop(self):
        """停止资源监控"""
        self._running = False
        
        if self._monitoring_task:
            self._monitoring_task.cancel()
        if self._cleanup_task:
            self._cleanup_task.cancel()
        
        logger.info("Resource Monitor stopped")
    
    async def _monitoring_loop(self):
        """监控循环"""
        while self._running:
            try:
                # 收集系统指标
                await self._collect_system_metrics()
                
                # 收集GPU指标
                await self._collect_gpu_metrics()
                
                # 检查告警规则
                await self._check_alert_rules()
                
                await asyncio.sleep(self.monitoring_interval)
                
            except Exception as e:
                logger.error(f"Error in monitoring loop: {e}")
                await asyncio.sleep(60)
    
    async def _collect_system_metrics(self):
        """收集系统指标"""
        try:
            # CPU使用率
            cpu_percent = psutil.cpu_percent(interval=1)
            
            # 内存使用情况
            memory = psutil.virtual_memory()
            
            # 磁盘使用情况
            disk = psutil.disk_usage('/')
            
            # 网络统计
            network = psutil.net_io_counters()
            
            # 负载平均值
            try:
                load_avg = psutil.getloadavg()
            except AttributeError:
                # Windows系统没有getloadavg
                load_avg = [0.0, 0.0, 0.0]
            
            # 进程数量
            process_count = len(psutil.pids())
            
            # 系统温度（如果可用）
            temperature = None
            try:
                temps = psutil.sensors_temperatures()
                if temps:
                    # 取第一个温度传感器的值
                    for name, entries in temps.items():
                        if entries:
                            temperature = entries[0].current
                            break
            except:
                pass
            
            # 创建指标对象
            metrics = SystemMetrics(
                timestamp=datetime.now(),
                cpu_percent=cpu_percent,
                memory_percent=memory.percent,
                memory_used=memory.used // (1024 * 1024),  # 转换为MB
                memory_total=memory.total // (1024 * 1024),
                disk_percent=disk.percent,
                disk_used=disk.used // (1024 * 1024 * 1024),  # 转换为GB
                disk_total=disk.total // (1024 * 1024 * 1024),
                network_sent=network.bytes_sent,
                network_recv=network.bytes_recv,
                load_average=list(load_avg),
                process_count=process_count,
                temperature=temperature
            )
            
            # 添加到历史记录
            self.system_metrics_history.append(metrics)
            
            # 限制历史记录大小
            if len(self.system_metrics_history) > self.max_history_points:
                self.system_metrics_history = self.system_metrics_history[-self.max_history_points:]
            
        except Exception as e:
            logger.error(f"Failed to collect system metrics: {e}")
    
    async def _collect_gpu_metrics(self):
        """收集GPU指标"""
        try:
            gpu_info = await gpu_resource_manager.get_gpu_info()
            
            for gpu_id, gpu in gpu_info.items():
                metrics = GPUMetrics(
                    timestamp=datetime.now(),
                    gpu_id=gpu_id,
                    utilization=gpu.utilization,
                    memory_used=gpu.memory_used,
                    memory_total=gpu.memory_total,
                    temperature=gpu.temperature,
                    power_usage=gpu.power_usage,
                    fan_speed=None  # 暂时不支持风扇速度
                )
                
                # 添加到历史记录
                if gpu_id not in self.gpu_metrics_history:
                    self.gpu_metrics_history[gpu_id] = []
                
                self.gpu_metrics_history[gpu_id].append(metrics)
                
                # 限制历史记录大小
                if len(self.gpu_metrics_history[gpu_id]) > self.max_history_points:
                    self.gpu_metrics_history[gpu_id] = self.gpu_metrics_history[gpu_id][-self.max_history_points:]
            
        except Exception as e:
            logger.error(f"Failed to collect GPU metrics: {e}")
    
    async def _check_alert_rules(self):
        """检查告警规则"""
        current_time = datetime.now()
        
        for rule_id, rule in self.alert_rules.items():
            if not rule.enabled:
                continue
            
            try:
                # 获取当前指标值
                current_value = await self._get_current_metric_value(rule.metric_type)
                if current_value is None:
                    continue
                
                # 检查阈值
                triggered = self._evaluate_threshold(current_value, rule.threshold, rule.operator)
                
                if triggered:
                    # 检查是否已经有活跃告警
                    if rule_id in self.active_alerts:
                        alert = self.active_alerts[rule_id]
                        # 检查持续时间
                        if (current_time - alert['start_time']).total_seconds() >= rule.duration:
                            # 更新告警
                            alert['last_update'] = current_time
                            alert['current_value'] = current_value
                            
                            # 触发告警通知
                            await self._trigger_alert(rule, current_value)
                    else:
                        # 创建新的活跃告警
                        self.active_alerts[rule_id] = {
                            'rule': rule,
                            'start_time': current_time,
                            'last_update': current_time,
                            'current_value': current_value
                        }
                else:
                    # 清除活跃告警
                    if rule_id in self.active_alerts:
                        del self.active_alerts[rule_id]
                        logger.info(f"Alert {rule.name} resolved")
                
            except Exception as e:
                logger.error(f"Error checking alert rule {rule_id}: {e}")
    
    async def _get_current_metric_value(self, metric_type: str) -> Optional[float]:
        """获取当前指标值"""
        try:
            if metric_type == "cpu_percent":
                if self.system_metrics_history:
                    return self.system_metrics_history[-1].cpu_percent
            
            elif metric_type == "memory_percent":
                if self.system_metrics_history:
                    return self.system_metrics_history[-1].memory_percent
            
            elif metric_type == "disk_percent":
                if self.system_metrics_history:
                    return self.system_metrics_history[-1].disk_percent
            
            elif metric_type == "gpu_utilization":
                # 返回所有GPU的平均利用率
                gpu_info = await gpu_resource_manager.get_gpu_info()
                if gpu_info:
                    utilizations = [gpu.utilization for gpu in gpu_info.values()]
                    return sum(utilizations) / len(utilizations)
            
            elif metric_type == "gpu_temperature":
                # 返回所有GPU的最高温度
                gpu_info = await gpu_resource_manager.get_gpu_info()
                if gpu_info:
                    temperatures = [gpu.temperature for gpu in gpu_info.values() if gpu.temperature is not None]
                    return max(temperatures) if temperatures else None
            
            return None
            
        except Exception as e:
            logger.error(f"Failed to get metric value for {metric_type}: {e}")
            return None
    
    def _evaluate_threshold(self, value: float, threshold: float, operator: str) -> bool:
        """评估阈值条件"""
        if operator == ">=":
            return value >= threshold
        elif operator == ">":
            return value > threshold
        elif operator == "<=":
            return value <= threshold
        elif operator == "<":
            return value < threshold
        elif operator == "==":
            return abs(value - threshold) < 0.001  # 浮点数比较
        else:
            return False
    
    async def _trigger_alert(self, rule: AlertRule, current_value: float):
        """触发告警"""
        # 检查是否在冷却期内
        if rule.last_triggered:
            cooldown_period = timedelta(minutes=15)  # 15分钟冷却期
            if datetime.now() - rule.last_triggered < cooldown_period:
                return
        
        # 更新最后触发时间
        rule.last_triggered = datetime.now()
        
        # 记录告警日志
        logger.warning(f"ALERT: {rule.name} - Current value: {current_value}, Threshold: {rule.threshold}")
        
        # 这里可以添加告警通知逻辑，如发送邮件、短信、Webhook等
        await self._send_alert_notification(rule, current_value)
    
    async def _send_alert_notification(self, rule: AlertRule, current_value: float):
        """发送告警通知"""
        # 构建告警消息
        message = {
            "alert_id": rule.rule_id,
            "alert_name": rule.name,
            "metric_type": rule.metric_type,
            "current_value": current_value,
            "threshold": rule.threshold,
            "operator": rule.operator,
            "timestamp": datetime.now().isoformat(),
            "severity": self._get_alert_severity(rule, current_value)
        }
        
        # 这里可以实现具体的通知逻辑
        # 例如：发送到消息队列、调用Webhook、发送邮件等
        logger.info(f"Alert notification: {json.dumps(message, indent=2)}")
    
    def _get_alert_severity(self, rule: AlertRule, current_value: float) -> str:
        """获取告警严重程度"""
        if rule.metric_type in ["gpu_temperature", "cpu_percent", "memory_percent"]:
            if current_value >= rule.threshold * 1.2:  # 超过阈值20%
                return "critical"
            elif current_value >= rule.threshold * 1.1:  # 超过阈值10%
                return "high"
            else:
                return "medium"
        else:
            return "medium"
    
    async def _cleanup_loop(self):
        """清理循环"""
        while self._running:
            try:
                await self._cleanup_old_metrics()
                await asyncio.sleep(3600)  # 每小时清理一次
            except Exception as e:
                logger.error(f"Error in cleanup loop: {e}")
                await asyncio.sleep(3600)
    
    async def _cleanup_old_metrics(self):
        """清理旧的指标数据"""
        cutoff_time = datetime.now() - timedelta(hours=self.history_retention_hours)
        
        # 清理系统指标
        self.system_metrics_history = [
            metrics for metrics in self.system_metrics_history
            if metrics.timestamp >= cutoff_time
        ]
        
        # 清理GPU指标
        for gpu_id in self.gpu_metrics_history:
            self.gpu_metrics_history[gpu_id] = [
                metrics for metrics in self.gpu_metrics_history[gpu_id]
                if metrics.timestamp >= cutoff_time
            ]
        
        logger.info("Cleaned up old metrics data")
    
    # 公共API方法
    
    async def get_current_system_metrics(self) -> Optional[SystemMetrics]:
        """获取当前系统指标"""
        return self.system_metrics_history[-1] if self.system_metrics_history else None
    
    async def get_system_metrics_history(self, hours: int = 1) -> List[SystemMetrics]:
        """获取系统指标历史"""
        cutoff_time = datetime.now() - timedelta(hours=hours)
        return [
            metrics for metrics in self.system_metrics_history
            if metrics.timestamp >= cutoff_time
        ]
    
    async def get_gpu_metrics_history(self, gpu_id: int, hours: int = 1) -> List[GPUMetrics]:
        """获取GPU指标历史"""
        if gpu_id not in self.gpu_metrics_history:
            return []
        
        cutoff_time = datetime.now() - timedelta(hours=hours)
        return [
            metrics for metrics in self.gpu_metrics_history[gpu_id]
            if metrics.timestamp >= cutoff_time
        ]
    
    async def get_resource_summary(self) -> Dict[str, Any]:
        """获取资源摘要"""
        current_system = await self.get_current_system_metrics()
        gpu_info = await gpu_resource_manager.get_gpu_info()
        cluster_info = await gpu_resource_manager.get_cluster_info()
        queue_status = await ai_task_scheduler.get_queue_status()
        
        return {
            "timestamp": datetime.now().isoformat(),
            "system": {
                "cpu_percent": current_system.cpu_percent if current_system else 0,
                "memory_percent": current_system.memory_percent if current_system else 0,
                "disk_percent": current_system.disk_percent if current_system else 0,
                "load_average": current_system.load_average if current_system else [0, 0, 0],
                "process_count": current_system.process_count if current_system else 0,
                "temperature": current_system.temperature if current_system else None
            },
            "gpu_cluster": {
                "total_gpus": cluster_info.total_gpus,
                "available_gpus": cluster_info.available_gpus,
                "average_utilization": cluster_info.average_utilization,
                "memory_usage_percent": (cluster_info.used_memory / max(cluster_info.total_memory, 1)) * 100
            },
            "task_queue": queue_status,
            "active_alerts": len(self.active_alerts),
            "alert_rules": len([rule for rule in self.alert_rules.values() if rule.enabled])
        }
    
    def add_alert_rule(self, rule: AlertRule):
        """添加告警规则"""
        self.alert_rules[rule.rule_id] = rule
        logger.info(f"Added alert rule: {rule.name}")
    
    def remove_alert_rule(self, rule_id: str) -> bool:
        """移除告警规则"""
        if rule_id in self.alert_rules:
            del self.alert_rules[rule_id]
            # 清除相关的活跃告警
            if rule_id in self.active_alerts:
                del self.active_alerts[rule_id]
            logger.info(f"Removed alert rule: {rule_id}")
            return True
        return False
    
    def get_alert_rules(self) -> List[AlertRule]:
        """获取所有告警规则"""
        return list(self.alert_rules.values())
    
    def get_active_alerts(self) -> List[Dict[str, Any]]:
        """获取活跃告警"""
        return [
            {
                "rule_id": rule_id,
                "rule_name": alert['rule'].name,
                "metric_type": alert['rule'].metric_type,
                "current_value": alert['current_value'],
                "threshold": alert['rule'].threshold,
                "start_time": alert['start_time'].isoformat(),
                "duration": (datetime.now() - alert['start_time']).total_seconds()
            }
            for rule_id, alert in self.active_alerts.items()
        ]


# 全局资源监控器实例
resource_monitor = ResourceMonitor()